import json
import linecache
import logging
import random

import gensim
import numpy as np
from gensim import corpora
from tqdm import tqdm

random.seed(123)

logger = logging.getLogger(__name__)
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)


def load_json(url):
    with open(url) as jsonFile:
        data = json.load(jsonFile)
        return data


def load_text(url):
    text = []
    file = linecache.getlines(url)
    for line in file:
        text.append(line.split())
    return text


# 得到纯字符集
'''
def stopwords(sent):
    # pattern = re.compile(u"[\u4e00-\u9fa5]+")
    new_sent = []
    for word in sent:
        if not word in stopword:
            new_sent.append(word)
    return new_sent
'''


def loss(input, target):
    return np.linalg.norm(input - target)


def load_model(url):
    return gensim.models.HdpModel.load(url, mmap='r')


def sent2vec(sent, dictionary, model, topicN):
    doc2bow = dictionary.doc2bow(sent)
    result = model[doc2bow]

    vec = [0.] * topicN
    for item in result:
        vec[item[0]] = item[1]
    return vec
    '''
    flag = 0
    for idx in range(150):
        if flag < len(result) and idx == result[flag][0]:
            vec.append(result[flag][1])
            flag += 1
        else:
            vec.append(0.0)

    return vec
    '''


def distance(corpus, model, testset, topicN):
    text = load_text(corpus)
    dictionary = corpora.Dictionary(text)
    model = load_model(model)
    test_set = load_json(testset)
    label = []

    pbar = tqdm(total=len(test_set))
    pbar.set_description('distance')
    pbar.update()

    for cqa in test_set:
        context = ' '.join(cqa['context']).split()
        query = cqa['query'].split()
        reply = cqa['reply']
        score = []

        for r in reply:
            sent = query + r.split()
            context_v = sent2vec(context, dictionary, model, topicN)
            sent_v = sent2vec(sent, dictionary, model, topicN)
            score.append(cos(context_v, sent_v))

        if max(score) == 0.0:
            label.append(random.randint(0, 9))
        else:
            label.append(score.index(max(score)))
        pbar.update(1)
    return label


def cos(x, y):
    a = np.array(x)
    b = np.array(y)
    return np.inner(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def get_topic():
    model = load_model('hdp_model')
    return len(model.get_topics())


def word2id(word):
    model = load_model('hdp_model')
    m = model.id2word.id2token
    word2id = {v: k for k, v in m.items()}
    return word2id[word.decode('utf-8')]


def write2file(str, path):
    with open(path, 'a') as file:
        file.writelines(str)
        file.writelines('\n')


if __name__ == '__main__':
    # stopword = load_json('data/stopwords.json')
    topicN = get_topic()
    labels = distance('out/corpus_no_stopwords.txt', 'hdp_model', 'raw/submit.json', topicN)
    for l in labels:
        write2file(str(l), 'out/hdp_submit.txt')
    # print(len(get_topic()))
